import sys
import faiss
import numpy as np

text_path = "../utils/embedded_files/embedded_user_query.txt"

def text_to_embedding(txt_path):
    args = []
    num = 0

    with open(txt_path, "r", encoding="utf-8") as doc:
        text = doc.read()
        paras = text.split("\n")

    for para in paras:
        if para != "":
            num = num + 1
            arg = float(para)
            args.append(arg)

    numpy_array = np.array(args, dtype=np.float64)
    numpy_2d_array = numpy_array.reshape(1, num)
    return numpy_2d_array


# 把query进行向量化，通过创建的index，获取最相似的top_k个结果
def data_recall(faiss_index, top_k):
    # 将查询文本编码为向量
    query_embedding = text_to_embedding(text_path)
    # 使用 FAISS 索引搜索
    distance, index = faiss_index.search(query_embedding, top_k)
    # 返回最相似结果的索引
    return index, distance

# 将最相似的索引转换为文本
def get_similar_texts(indices, texts):
    similar_texts = [texts[idx] for idx in indices[0]]
    return similar_texts

def get_file_name(line):
    if line < 18:
        return "01_银河麒麟桌面操作系统 V101产品白皮书"
    elif 18 <= line < 36:
        return "02_银河麒麟桌面操作系统V10 SP1 2503常见问题(FAQ)"
    elif 36 <= line < 66:
        return "03_银河麒麟桌面操作系统V10 SP1 2503产品安装手册"
    elif 66 <= line < 167:
        return "04_银河麒麟桌面操作系统V10 SP1 2503产品用户手册"


if __name__ == '__main__':

    # index地址
    index_path = "../utils/index/RAG.index"
    index = faiss.read_index(index_path)

    headers_index_path = "../utils/index/header/header_RAG.index"
    headers_index = faiss.read_index(headers_index_path)

    txt_path = "../utils/text.txt"
    with open(txt_path, "r", encoding="utf-8") as doc:
        text = doc.read()
        chunks = text.split("\n")

    top_k1= 2

    top_k2 = 3

    header_indices, header_distances = data_recall(headers_index, top_k1)

    # print(header_indices)
    # print("---")
    # print(header_distances)
    # print("---")

    similar_header_texts = get_similar_texts(header_indices, chunks)

    for i in range(top_k1):
        file_name = get_file_name(header_indices[0][i])
        print(file_name)
        print(similar_header_texts[i])
        print("---")

    # print("\n-------------------------------------------------------\n")

    indices, distances = data_recall(index, top_k2)

    # print(indices)
    # print("---")
    # # print(distances)
    # print("---")

    similar_texts = get_similar_texts(indices, chunks)

    for i in range(top_k2):
        file_name = get_file_name(indices[0][i])
        print(file_name)
        print(similar_texts[i])
        print("---")